# Description: This script calculates the codon adaptation index (CAI) for a given mRNA sequence.
# Author: Shibo Li, MiQroEra Inc.
# Date: 2023-08-21
# License: MIT License
# Version: 1.0

import numpy as np
from modules.loadCodonUsage import load_codon_usage_table


def calculate_CAI(mrna_seq, codon_table):
    """
    The function is to calculate the codon adaptation index (CAI) for a given mRNA sequence.
    :param mrna_seq: mRNA sequence
    :param codon_table: codon usage table
    :return: CAI value
    :rtype: float
    """
    mrna_seq = mrna_seq.upper()
    codon_count = len(mrna_seq) // 3

    if len(mrna_seq) % 3 != 0:
        raise ValueError(
            "The length of the mRNA sequence should be a multiple of 3.")

    for i in range(0, len(mrna_seq), 3):
        codon = mrna_seq[i:i+3]
        if codon not in [cod for sublist in codon_table.values() for cod in sublist]:
            raise ValueError(
                f"Invalid codon '{codon}' found in the mRNA sequence.")

    codon_weight = {}
    for aa in codon_table:
        max_freq = max(codon_table[aa].values())
        for codon in codon_table[aa]:
            codon_weight[codon] = codon_table[aa][codon] / max_freq

    cai = np.prod([codon_weight[mrna_seq[i:i+3]]
                  for i in range(0, len(mrna_seq), 3)]) ** (1/codon_count)
    return cai


if __name__ == "__main__":
    filepath = "data/codon_usage/Caenorhabditis_elegans.csv"
    codon_table = load_codon_usage_table(filepath)
    test_sequence = "AUGGCGGGGCGCCGAGCCCAGACGGGCUCGGCGCCCCCCAGGCCCGCGGCACCGCACCCGCGGCCGGCAAGUAGGGCAUUCCCGCAGCAUUGCCGGCCCCGGGACGCCGAACGGCCGCCGAGCCCGCGCUCCCCGCUCAUGCCUGGGUGCGAGCUGCCGGUGGGCACAUGUCCCGACAUGUGCCCAGCGGCAGAGCGCGCCCAGCGUGAGCGGGAGCACCGGCUCCACCGGCUGGAGGUCGUCCCGGGGUGUCGGCAGGACCCACCGCGCGCUGAUCCUCAGCGCGCGGUGAAGGAGUAUUCCCGGCCGGCGGCCGGGAAGCCUCGGCCGCCACCGAGCCAGCUAAGACCGCCCUCGGUCUUGCUGGCUACGGUGCGGUACCUGGCCGGGGAGGUGGCCGAGUCGGCCGACAUAGCGCGAGCAGAGGUGGCCAGCUUCGUGGCCGACAGGCUACGGGCGGUUCUGCUGGACCUGGCCCUGCAGGGGGCCGGGGACGCAGAAGCGGCCGUAGUCCUGGAGGCCGCGCUGGCCACCCUGCUCACGGUAGUCGCCCGACUCGGCCCCGACGCGGCCAGGGGUCCUGCCGACCCGGUAUUGCUGCAAGCGCAGGUGCAGGAGGGUUUCGGGUCCCUGCGCCGGUGCUAUGCACGCGGCGCAGGGCCCCACCCUCGGCAGCCUGCGUUUCAGGGCCUCUUCCUCCUCUACAACCUGGGGAGCGUGGAGGCCCUGCACGAAGUCUUGCAGCUUCCCGCUGCCCUGAGGGCCUGUCCACCGCUGCGCAAAGCGCUAGCGGUGGACGCGGCCUUCAGGGAGGGGAACGCUGCAAGACUUUUCCGUCUGCUGCAGACGCUGCCCUACUUGCCGAGCUGCGCGGUGCAGUGCCACGUGGGCCACGCGCGUCGAGAAGCGCUCGCGCGUUUCGCGCGAGCCUUCUCGACGCCCAAGGGGCAGACCCUGCCCCUUGGGUUCAUGGUGAACCUACUCGCUCUUGACGGUCUGAGGGAAGCCCGUGACCUGUGUCAGGCUCACGGGCUUCCCUUAGACGGUGAAGAGCGAGUAGUGUUCCUACGCGGGCGCUACGUGGAGGAGGGCCUUCCUCCUGCCUCCACCUGUAAGGUCCUCGUCGAGAGUAAGCUCAGAGGGCGGACUCUCGAGGAGGUCGUUAUGGCGGAGGAGGAGGACGAGGGCACCGACCGGCCAGGUAGCCCCGCG"
    print(
        f"CAI for test sequence: {calculate_CAI(test_sequence, codon_table)}")
